#include "kwdb/batcher.h"
#include <iostream>
#include <cstdlib>

namespace kwdb {

BatcherConfig BatcherConfig::FromEnvironment() {
    BatcherConfig config;

    const char* batch_size_env = getenv("KWDB_BATCH_SIZE");
    if (batch_size_env) {
        config.max_batch_size = std::stoul(batch_size_env);
    }

    const char* linger_env = getenv("KWDB_FLUSH_INTERVAL_MS");
    if (linger_env) {
        config.linger_ms = std::stoi(linger_env);
    }

    const char* pending_env = getenv("KWDB_MAX_PENDING_BATCHES");
    if (pending_env) {
        config.max_pending_batches = std::stoul(pending_env);
    }

    return config;
}

Batcher::Batcher(const BatcherConfig& config, BatchCallback callback)
    : config_(config), callback_(callback) {

    current_batch_.reserve(config_.max_batch_size);
    last_add_time_ = std::chrono::steady_clock::now();

    // Start flush thread
    flush_thread_ = std::thread(&Batcher::FlushThread, this);
}

Batcher::~Batcher() {
    Close();
}

Status Batcher::Add(const Point& point) {
    std::unique_lock<std::mutex> lock(current_mutex_);

    current_batch_.push_back(point);
    last_add_time_ = std::chrono::steady_clock::now();

    // Check if batch is full
    if (current_batch_.size() >= config_.max_batch_size) {
        // Move batch to pending queue
        std::vector<Point> batch = std::move(current_batch_);
        current_batch_.clear();
        current_batch_.reserve(config_.max_batch_size);

        lock.unlock();

        // Add to pending queue (may block on backpressure)
        {
            std::unique_lock<std::mutex> queue_lock(queue_mutex_);

            // Wait if too many pending batches (backpressure)
            backpressure_cv_.wait(queue_lock, [this] {
                return pending_batches_.size() < config_.max_pending_batches || !running_;
            });

            if (!running_) {
                return Status(-1, "Batcher is closed");
            }

            pending_batches_.push(std::move(batch));
            pending_count_ = pending_batches_.size();
        }

        queue_cv_.notify_one();
    }

    return Status();
}

Status Batcher::Add(const std::vector<Point>& points) {
    Status status;
    for (const auto& point : points) {
        status = Add(point);
        if (!status.ok()) {
            return status;
        }
    }
    return Status();
}

Status Batcher::Flush() {
    // Flush current batch first
    {
        std::unique_lock<std::mutex> lock(current_mutex_);
        if (!current_batch_.empty()) {
            std::vector<Point> batch = std::move(current_batch_);
            current_batch_.clear();
            current_batch_.reserve(config_.max_batch_size);

            lock.unlock();

            // Add to pending queue
            {
                std::unique_lock<std::mutex> queue_lock(queue_mutex_);
                pending_batches_.push(std::move(batch));
                pending_count_ = pending_batches_.size();
            }
            queue_cv_.notify_one();
        }
    }

    // Wait for all pending batches to be flushed
    {
        std::unique_lock<std::mutex> lock(queue_mutex_);
        backpressure_cv_.wait(lock, [this] {
            return pending_batches_.empty();
        });
    }

    return Status();
}

void Batcher::Close() {
    if (running_.exchange(false)) {
        // Flush remaining data
        Flush();

        // Wake up flush thread to exit
        queue_cv_.notify_all();

        // Wait for flush thread
        if (flush_thread_.joinable()) {
            flush_thread_.join();
        }
    }
}

void Batcher::FlushThread() {
    while (running_) {
        std::unique_lock<std::mutex> lock(queue_mutex_);

        // Wait for data or timeout
        auto timeout_point = std::chrono::steady_clock::now() +
                           std::chrono::milliseconds(config_.linger_ms);

        queue_cv_.wait_until(lock, timeout_point, [this] {
            return !pending_batches_.empty() || !running_;
        });

        // Check for linger timeout on current batch
        if (pending_batches_.empty() && running_) {
            lock.unlock();

            std::unique_lock<std::mutex> current_lock(current_mutex_);
            auto now = std::chrono::steady_clock::now();
            auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                now - last_add_time_).count();

            if (!current_batch_.empty() && elapsed >= config_.linger_ms) {
                // Timeout reached, flush current batch
                std::vector<Point> batch = std::move(current_batch_);
                current_batch_.clear();
                current_batch_.reserve(config_.max_batch_size);

                current_lock.unlock();

                // Process the batch
                FlushBatch(std::move(batch));
            }
            continue;
        }

        // Process pending batches
        while (!pending_batches_.empty() && running_) {
            std::vector<Point> batch = std::move(pending_batches_.front());
            pending_batches_.pop();
            pending_count_ = pending_batches_.size();

            lock.unlock();

            // Notify waiting producers (backpressure relief)
            backpressure_cv_.notify_all();

            // Process the batch
            FlushBatch(std::move(batch));

            lock.lock();
        }
    }

    // Final flush on shutdown
    std::unique_lock<std::mutex> lock(queue_mutex_);
    while (!pending_batches_.empty()) {
        std::vector<Point> batch = std::move(pending_batches_.front());
        pending_batches_.pop();

        lock.unlock();
        FlushBatch(std::move(batch));
        lock.lock();
    }
}

Status Batcher::FlushBatch(std::vector<Point>&& batch) {
    if (batch.empty()) {
        return Status();
    }

    Status status = callback_(batch);

    if (status.ok()) {
        batches_flushed_++;
        std::cout << "[Batcher] Flushed batch of " << batch.size()
                  << " points, total batches: " << batches_flushed_ << std::endl;
    } else {
        std::cerr << "[Batcher] Failed to flush batch: " << status.message << std::endl;
    }

    return status;
}

} // namespace kwdb